
#pragma once

#include <ttg/services.h>

#include "Common/ICalculationCore.h"
#include "PairedBodies.h"
#include "cudaBindings.h"

using namespace ttg::pipeline;
using namespace ttg::utils;

class ttgBodySplitter1 :public Node<InterleavedBodies::SubVector, InterleavedBodies::SubVector>
{
	private:
		SubPipeline<PairedBodies, PairedBodies> spInteractor;		//reference to sub-pipeline
	protected:
		void onProcess(std::vector<PairedBodies> &res)
		{
			sendThis(res[0].first);
			sendThis(res[0].second);
		}
		virtual void process(InterleavedBodies::SubVector data)
		{
			int threshold = getContext().getVariable<int>("threshold");
			if ((ttg::math::int64)data.getSize() * data.getSize() <= 2 * threshold)
			{ sendNext(data); }
			else
			{
				InterleavedBodies::SubVector v1(data, 0, data.getSize() / 2);
				InterleavedBodies::SubVector v2(data, data.getSize() / 2, data.getSize());
				sendTo(spInteractor, PairedBodies(v1, v2), make_delegate(this, &ttgBodySplitter1::onProcess));
			}
		}
		virtual void onModification()
		{
			spInteractor = createSubPipeline<PairedBodies, PairedBodies>("spInteractor");
		}
	public:
		inline ttgBodySplitter1()
			:Node("ttgBodySplitter1")
		{ /*nothing*/ }
};

class ttgBodyInteractor1 :public Node<InterleavedBodies::SubVector, InterleavedBodies::SubVector>
{
	public:
		virtual void process(InterleavedBodies::SubVector data)
		{
			for (size_t i = 0; i < data.getSize(); i++)
				for (size_t j = i + 1; j < data.getSize(); j++)
					data.interactWithEachOther(i, j);
			sendNext(data);
		}
		inline ttgBodyInteractor1()
			:Node("ttgBodyInteractor1")
		{ addBinding(new cudaBodyInteractor1()); }
};

class ttgBodyUpdater :public Node<InterleavedBodies::SubVector, InterleavedBodies::SubVector>
{
	public:
		virtual void process(InterleavedBodies::SubVector data)
		{
			float dt = getContext().getVariable<float>("dt");
			for (size_t i = 0; i < data.getSize(); i++)
				data.update(i, dt);
			sendNext(data);
		}
		inline ttgBodyUpdater()
			:Node("ttgBodyUpdater")
		{ /*nothing*/ }
};

class ttgBodySplitter2 :public Node<PairedBodies, PairedBodies>
{
	private:
		SubPipeline<PairedBodies, PairedBodies> spInteractor;		//reference to sub-pipeline.
	protected:
		//Recrosses sub-vectors and sends them again.
		void recrossSubVectors(std::vector<PairedBodies> &data)
		{
			std::vector<PairedBodies> pairs;
			pairs.push_back(PairedBodies(data[0].first, data[1].second));
			pairs.push_back(PairedBodies(data[1].first, data[0].second));
			sendTo(spInteractor, pairs);
		}
		//Node<...> member.
		virtual void process(PairedBodies data)
		{
			InterleavedBodies::SubVector &v1 = data.first;
			InterleavedBodies::SubVector &v2 = data.second;
			int threshold = getContext().getVariable<int>("threshold");
			if ((ttg::math::int64)v1.getSize() * v2.getSize() < threshold)
			{ sendNext(data); }
			else
			{
				//Preparing.
				InterleavedBodies::SubVector v11(v1, 0, v1.getSize() / 2);
				InterleavedBodies::SubVector v12(v1, v1.getSize() / 2, v1.getSize());
				InterleavedBodies::SubVector v21(v2, 0, v2.getSize() / 2);
				InterleavedBodies::SubVector v22(v2, v2.getSize() / 2, v2.getSize());
				std::vector<PairedBodies> pairs;
				//Sending direct comparisions.
				pairs.push_back(PairedBodies(v11, v21));
				pairs.push_back(PairedBodies(v12, v22));
				sendTo(spInteractor, pairs, make_delegate(this, &ttgBodySplitter2::recrossSubVectors));
			}
		}
		//Node<...> member.
		virtual void onModification()
		{
			spInteractor = createSubPipeline<PairedBodies, PairedBodies>("spInteractor");
		}
	public:
		inline ttgBodySplitter2()
			:Node("ttgBodySplitter2")
		{ /*nothing*/ }
};

class ttgBodyInteractor2 :public Node<PairedBodies, PairedBodies>
{
	public:
		virtual void process(PairedBodies data)
		{
			InterleavedBodies::SubVector &v1 = data.first;
			InterleavedBodies::SubVector &v2 = data.second;
			for (size_t i = 0; i < v1.getSize(); i++)
				for (size_t j = 0; j < v2.getSize(); j++)
					v1.interactWithEachOther(i, v2, j);
			sendNext(data);
		}

		inline ttgBodyInteractor2()
			:Node("ttgBodyInteractor2")
		{ addBinding(new cudaBodyInteractor2()); }
};

//Hack to avoid usage of unimplemented function.
//Will be replaced by obvious code (in future).
class ttgDirtyHack :public ContextedNode<PairedBodies,
										 PairedBodies>
{
	private:
		ttg::threads::FastCriticalSection locker;
		ContextMap<std::vector<PairedBodies>> incoming_data;
	public:
		virtual void onBegin(ContextID ID)
		{ incoming_data[ID].clear(); }

		virtual void onEnd(ContextID ID)
		{
			if (incoming_data[ID].size() == 0)
				sendNext(getContext().getInputData<PairedBodies>(ID));
			else
				sendNext(incoming_data[ID]);
		}
		virtual void process(ContextID ID, PairedBodies data)
		{
			locker.enter();
			incoming_data[getContextID()].push_back(data);
			locker.leave();
		}
		inline ttgDirtyHack()
			:ContextedNode("ttgDirtyHack")
		{ /*nothing*/ }
};

class ttgCore_Advanced :public ICalculationCore
{
	private:
		Pipeline<InterleavedBodies::SubVector, InterleavedBodies::SubVector> *pl;
		ParameterSet *ttglib_advanced_filter;
		ParameterEnum *threshold;
		std::vector<int> threshold_values;
	public:
		inline ttgCore_Advanced()
		{
			//Creating pipeline with sub-pipeline.
			pl = new Pipeline<InterleavedBodies::SubVector, InterleavedBodies::SubVector>(
				wrap(new ttgBodySplitter1()) +
				wrap(new ttgBodyInteractor1()) +
				wrap(new ttgBodyUpdater()),
				"ttgLib Advanced Pipeline");
			pl->addSubPipeline(wrap(new ttgBodySplitter2()) +
							   wrap(new ttgBodyInteractor2()) +
							   wrap(new ttgDirtyHack()), "spInteractor");
			//Creating parameters.
			ttglib_advanced_filter = new ParameterSet("ttgLib(Advanced) settings");
			threshold = new ParameterEnum("threshold", ttglib_advanced_filter,
										  "64|128|256|512|1024|2048|4096");
			threshold_values.push_back(64);
			threshold_values.push_back(128);
			threshold_values.push_back(256);
			threshold_values.push_back(512);
			threshold_values.push_back(1024);
			threshold_values.push_back(2048);
			threshold_values.push_back(4096);
		}

		virtual double simulate(InterleavedBodies::SubVector bodies, float dt)
		{
			//Preparing context.
			UserContext context;
			context.registerVariable("dt", &dt);
			context.createVariable<int>("threshold") = threshold_values[*threshold] * 1024;
			pl->setUserContext(context);
			//Launching.
			double t_start = ttg::services::Timer::getRef()->getSeconds();
			pl->run(bodies);
			double t_end = ttg::services::Timer::getRef()->getSeconds();
			//Returning elapsed time.
			return t_end - t_start;
		}

		virtual ~ttgCore_Advanced()
		{
			SAFE_DELETE(pl);
			SAFE_DELETE(ttglib_advanced_filter);
			SAFE_DELETE(threshold);
		}
};
